deepmd/pt/loss/__init__.py (1)

25-31: Consider alphabetically ordering the __all__ list.

While the current ordering works functionally, maintaining an alphabetical order would improve readability and make future modifications easier to manage.

Here's the suggested reordering:

 __all__ = [
-    "DOSLoss",
     "DenoiseLoss",
+    "DOSLoss",
     "EnergySpinLoss",
     "EnergyStdLoss",
     "PropertyLoss",
     "TaskLoss",
     "TensorLoss",
 ]

deepmd/jax/env.py (1)

21-21: Consider adding documentation about jax_export usage.

Since jax_export is now part of the public API, it would be helpful to add a comment explaining its purpose and intended usage within the DeepMD-kit context, especially given its role in the DPA-2 support implementation.

 __all__ = [
     "jax",
+    # Exposed for model export functionality used in DPA-2 implementation
     "jax_export",
     "jnp",
     "nnx",
 ]

deepmd/pt/model/task/dos.py (1)

Inconsistent version compatibility checks detected across fitting classes

Most fitting classes use version 3, but deepmd/pt/model/task/polarizability.py uses version 4. This inconsistency needs to be addressed.

• deepmd/pt/model/task/polarizability.py: Uses version 4
• All other fitting classes: Use version 3
  • deepmd/pd/model/task/ener.py
  • deepmd/pd/model/task/invar_fitting.py
  • deepmd/pt/model/task/property.py
  • deepmd/pt/model/task/invar_fitting.py
  • deepmd/pt/model/task/ener.py
  • deepmd/pt/model/task/dipole.py
  • deepmd/pt/model/task/dos.py

🔗 Analysis chain

Line range hint 104-113: Verify version compatibility across fitting classes

The version compatibility check has been updated from 2 to 3, which is a breaking change.

Let's verify this change is consistent across other fitting classes:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Check version compatibility in other fitting classes
# Expected: All fitting classes should use version 3 for consistency

# Search for version compatibility checks in fitting classes
rg -A 1 "check_version_compatibility" --glob "deepmd/**/fitting*.py" --glob "deepmd/**/task/*.py"

Length of output: 2289

deepmd/dpmodel/model/__init__.py (1)

Line range hint 1-37: Consider adding a changelog entry

Since this change modifies the public API by exposing DPModelCommon, it would be helpful to document this change in the changelog to help users understand the implications of this addition.

deepmd/dpmodel/fitting/property_fitting.py (1)

81-81: Add documentation for the new dim_case_embd parameter.

The new parameter has been added correctly, but its documentation is missing in the class docstring. Please update the docstring to include:

• Purpose and usage of dim_case_embd
• Expected value range
• Default behavior when set to 0

Also applies to: 103-103

deepmd/pt/model/atomic_model/pairtab_atomic_model.py (2)

144-151: Enhance docstring to indicate unsupported operation

The implementation correctly raises NotImplementedError for unsupported case identification. Consider updating the docstring to explicitly mention that this operation is not supported by PairTabAtomicModel.

     def set_case_embd(self, case_idx: int):
         """
-        Set the case embedding of this atomic model by the given case_idx,
-        typically concatenated with the output of the descriptor and fed into the fitting net.
+        Set the case embedding of this atomic model by the given case_idx.
+        
+        Note: This operation is not supported by PairTabAtomicModel as it does not
+        utilize case identification in its implementation.
+        
+        Parameters
+        ----------
+        case_idx : int
+            The index of the case to be embedded.
+            
+        Raises
+        ------
+        NotImplementedError
+            Always raised as case identification is not supported.
         """

---

497-506: Enhance docstring with parameter documentation

The implementation correctly indicates that compression is not supported. Consider enhancing the docstring to document the parameters and explain why compression isn't supported for this model type.

     def enable_compression(
         self,
         min_nbor_dist: float,
         table_extrapolate: float = 5,
         table_stride_1: float = 0.01,
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Pairtab model does not support compression."""
+        """
+        Enable model compression (not supported).
+        
+        PairTabAtomicModel does not support compression as it already uses a
+        pre-computed tabulation approach for efficiency.
+        
+        Parameters
+        ----------
+        min_nbor_dist : float
+            Minimum neighbor distance for compression (not used).
+        table_extrapolate : float, optional
+            Table extrapolation parameter (not used), by default 5.
+        table_stride_1 : float, optional
+            First table stride parameter (not used), by default 0.01.
+        table_stride_2 : float, optional
+            Second table stride parameter (not used), by default 0.1.
+        check_frequency : int, optional
+            Frequency of compression checks (not used), by default -1.
+        """

deepmd/env.py (1)

17-21: Consider adding docstring to document the public interface.

Since these variables are now part of the public interface, it would be helpful to add a module-level docstring explaining their purpose and usage.

Example addition:

+ """
+ Global configuration module for DeepMD-kit.
+
+ Exports:
+     GLOBAL_ENER_FLOAT_PRECISION: NumPy dtype for energy calculations
+     GLOBAL_NP_FLOAT_PRECISION: NumPy dtype for general calculations
+     SHARED_LIB_MODULE: Name of the shared library module
+     global_float_prec: String identifier for precision level ("float" or "double")
+ """

deepmd/pt/model/atomic_model/__init__.py (1)

50-56: LGTM! Consider grouping related model classes together.

The changes to __all__ maintain consistency with the imported classes. However, consider grouping related model classes together for better organization. For example:

• Base classes (BaseAtomicModel)
• Core models (DPAtomicModel)
• Property models (DPPropertyAtomicModel, DPDipoleAtomicModel)
• Energy models (DPEnergyAtomicModel, LinearEnergyAtomicModel, DPZBLLinearEnergyAtomicModel)
• Special models (PairTabAtomicModel)

deepmd/pt/model/task/dipole.py (1)

54-55: Enhance documentation for dim_case_embd parameter

While the parameter is properly documented, it would be helpful to provide more context about what case-specific embedding means and how it affects the dipole fitting network's behavior.

deepmd/pt/cxx_op.py (1)

Line range hint 1-100: Consider unifying error handling across PyTorch and Paddle backends.

Since this PR is implementing DPA-2 support in the Paddle backend, it would be beneficial to ensure consistent error handling strategies between PyTorch and Paddle implementations. Consider:

1. Implementing similar version and ABI compatibility checks in the Paddle backend
2. Using a shared error message formatting strategy across backends
3. Creating common utility functions for version compatibility checks

This will help maintain consistency and reduce duplication as the codebase grows with multiple backend support.

🧰 Tools

🪛 Ruff (0.8.0)

50-50: Yoda condition detected

Rewrite as pt_cxx11_abi_flag != PT_CXX11_ABI_FLAG

(SIM300)

deepmd/pd/model/descriptor/se_t_tebd.py (3)

386-387: Remove unused variable env_mat

The env_mat variable is assigned but never used in the deserialize method.

- env_mat = data.pop("env_mat")
+ data.pop("env_mat")  # Remove unused env_mat data

🧰 Tools

🪛 Ruff (0.8.0)

387-387: Local variable env_mat is assigned to but never used

Remove assignment to unused variable env_mat

(F841)

🪛 GitHub Check: CodeQL

[notice] 386-386: Unused local variable
Variable env_mat is not used.

---

823-827: Optimize tensor operations using einsum

The current implementation uses manual tensor operations that could be simplified and potentially optimized using paddle.einsum.

Consider replacing the manual tensor operations with einsum:

- env_ij = (
-     rr_i.unsqueeze(2) * rr_j.unsqueeze(1)
- ).sum(-1)
+ env_ij = paddle.einsum("ijm,ikm->ijk", rr_i, rr_j)

- res_ij = (
-     env_ij.unsqueeze(-1) * gg
- ).sum([1, 2])
+ res_ij = paddle.einsum("ijk,ijkm->im", env_ij, gg)

This change would make the code more readable and potentially more efficient as einsum operations are often optimized at a lower level.

Also applies to: 909-913

---

709-713: Simplify conditional assignment using ternary operator

The current if-else block can be simplified using a ternary operator for better readability.

- if callable(merged):
-     # only get data for once
-     sampled = merged()
- else:
-     sampled = merged
+ # only get data for once if callable
+ sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.0)

709-713: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

(SIM108)

deepmd/pd/utils/exclude_mask.py (1)

104-104: Remove or clarify the uncertain comment

The comment "may have a better place for this method..." suggests uncertainty about the method placement. If there's a better location for this method, it should be moved there. Otherwise, the comment should be removed to avoid confusion.

deepmd/pd/model/descriptor/repformers.py (3)

441-499: Document or remove commented-out code

The commented-out code block for spin handling appears to be a work in progress. If this is planned functionality, it should be properly documented with a TODO comment explaining the implementation plan. If not needed, it should be removed.

Consider adding a TODO comment at the start of the block:

# TODO: Implement spin handling functionality
# This block will handle:
# 1. Padding for spin components
# 2. Communication between processes
# 3. Virtual particle handling

---

365-522: Consider breaking down the forward method

The forward method is quite complex with multiple responsibilities. Consider extracting some logic into helper methods for better maintainability:

1. Neighbor list processing
2. Environment matrix computation
3. Layer processing
4. Rotation matrix computation

Example refactor:

def _process_neighbor_list(self, nlist, extended_atype):
    # Extract neighbor list processing logic

def _compute_environment_matrix(self, extended_coord, nlist, atype):
    # Extract environment matrix computation logic

def _process_layers(self, g1, g2, h2, nlist, nlist_mask, sw):
    # Extract layer processing logic

🧰 Tools

🪛 Ruff (0.8.0)

412-412: Local variable ng1 is assigned to but never used

Remove assignment to unused variable ng1

(F841)

---

435-435: Loop control variable idx not used within loop body

Rename unused idx to _idx

(B007)

🪛 GitHub Check: CodeQL

[notice] 416-416: Unused local variable
Variable ng1 is not used.

---

548-552: Simplify conditional assignment using ternary operator

The if-else block can be simplified using a ternary operator for better readability.

-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
+            sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.0)

548-552: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

(SIM108)

deepmd/dpmodel/fitting/dipole_fitting.py (2)

98-98: Add documentation for the new parameter dim_case_embd

The new parameter needs to be documented in the class docstring under the Parameters section. This will help users understand its purpose and usage.

Add this to the class docstring:

    Parameters
    ----------
+   dim_case_embd : int, optional
+       Dimension for case-specific embeddings. Defaults to 0.

---

Line range hint 1-234: Architecture Review: Consider performance optimization for large-scale computations

The implementation is well-structured and maintains type safety. However, given this is part of the DPA-2 support in Paddle backend, consider these optimizations:

1. The matrix multiplication in the call method could potentially benefit from Paddle's optimized operations
2. The reshape operations might be optimizable for better memory efficiency

Would you like assistance in exploring these optimizations?

deepmd/pd/model/network/init.py (2)

56-56: Rename ambiguous variable l for better readability

The single-letter variable l can be easily confused with the number 1 or uppercase I. Consider renaming it to something more descriptive like lower_bound.

-        l = norm_cdf((a - mean) / std)
-        u = norm_cdf((b - mean) / std)
+        lower_bound = norm_cdf((a - mean) / std)
+        upper_bound = norm_cdf((b - mean) / std)

-        tensor.uniform_(2 * l - 1, 2 * u - 1)
+        tensor.uniform_(2 * lower_bound - 1, 2 * upper_bound - 1)

🧰 Tools

🪛 Ruff (0.8.0)

56-56: Ambiguous variable name: l

(E741)

---

328-328: Add explicit stacklevel to warning messages

The warnings should include an explicit stacklevel to ensure they point to the correct location in the call stack.

-        warnings.warn("Initializing zero-element tensors is a no-op")
+        warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)

Also applies to: 378-378

🧰 Tools

🪛 Ruff (0.8.0)

328-328: No explicit stacklevel keyword argument found

(B028)

deepmd/pt/model/task/ener.py (1)

53-53: LGTM! Consider adding parameter documentation.

The addition of dim_case_embd parameter and its propagation to the superclass is well-implemented and maintains backward compatibility. However, the parameter's purpose and usage should be documented in the docstring.

Add parameter documentation:

    def __init__(
        self,
        ntypes: int,
        dim_descrpt: int,
        neuron: list[int] = [128, 128, 128],
        bias_atom_e: Optional[torch.Tensor] = None,
        resnet_dt: bool = True,
        numb_fparam: int = 0,
        numb_aparam: int = 0,
        dim_case_embd: int = 0,
        activation_function: str = "tanh",
        precision: str = DEFAULT_PRECISION,
        mixed_types: bool = True,
        seed: Optional[Union[int, list[int]]] = None,
        type_map: Optional[list[str]] = None,
        **kwargs,
    ) -> None:
+       """Initialize EnergyFittingNet.
+       
+       Args:
+           dim_case_embd: Dimension of case embedding. Defaults to 0.
+           [document other parameters...]
+       """

Also applies to: 71-71

deepmd/pt/model/descriptor/dpa2.py (1)

Line range hint 445-519: Add validation for the new type map.

While the implementation is thorough, consider adding validation for the new type map to ensure it's well-formed:

1. Check for duplicate types
2. Validate type names (e.g., non-empty strings)
3. Handle empty type maps gracefully

 def change_type_map(
     self, type_map: list[str], model_with_new_type_stat=None
 ) -> None:
     assert (
         self.type_map is not None
     ), "'type_map' must be defined when performing type changing!"
+    # Validate new type map
+    if not type_map:
+        raise ValueError("Type map cannot be empty")
+    if len(set(type_map)) != len(type_map):
+        raise ValueError("Duplicate types found in type map")
+    if any(not isinstance(t, str) or not t.strip() for t in type_map):
+        raise ValueError("Invalid type names in type map")
     remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map)

deepmd/dpmodel/fitting/polarizability_fitting.py (1)

Line range hint 1-324: Consider adding tests for case embedding functionality

While the implementation looks good, we should ensure proper test coverage for the new case embedding functionality in the polarizability fitting.

Would you like me to help generate test cases for:

1. Serialization/deserialization with case embedding
2. Forward pass with non-zero case embedding dimension
3. Backward compatibility with zero case embedding dimension

deepmd/infer/model_devi.py (2)

Line range hint 19-42: Consider simplifying overload signatures

The overload signatures use ellipsis (...) for the atomic parameter default values, but they could be more explicit.

Consider this improvement:

 @overload
 def calc_model_devi_f(
     fs: np.ndarray,
     real_f: Optional[np.ndarray] = None,
     relative: Optional[float] = None,
-    atomic: Literal[False] = ...,
+    atomic: Literal[False] = False,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ...

 @overload
 def calc_model_devi_f(
     fs: np.ndarray,
     real_f: Optional[np.ndarray] = None,
     relative: Optional[float] = None,
-    atomic: Literal[True] = ...,
+    atomic: Literal[True] = True,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ...

---

Based on the gathered evidence, I can now provide a final response:

Model deviation functionality needs to be implemented in Paddle backend

The verification reveals:

• DPA-2 descriptor is properly implemented in Paddle backend (deepmd/pd/model/descriptor/dpa2.py)
• However, model deviation functionality is missing in Paddle backend:
  • No model deviation implementation found in Paddle's inference module
  • No model deviation tests exist for Paddle backend
  • The functionality exists only in TensorFlow backend (deepmd/tf/infer/model_devi.py)

🔗 Analysis chain

Line range hint 1-450: Verify DPA-2 implementation completeness

While the code quality improvements are good, I don't see direct implementation of DPA-2 support as mentioned in the PR objectives.

Let's verify if DPA-2 related changes are present in other files:

---

The search results show that DPA-2 implementation exists in both Paddle and PyTorch backends:

• deepmd/pd/model/descriptor/dpa2.py
• deepmd/pt/model/descriptor/dpa2.py
• deepmd/jax/descriptor/dpa2.py

Let me verify if the model deviation functionality is properly implemented in the Paddle backend:

---

Let me check the model deviation implementation in more detail:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for DPA-2 related implementations
echo "Searching for DPA-2 related files..."
rg -l "dpa2|DPA2|DPA-2" --type py

echo "Searching for Paddle backend changes..."
rg -l "paddle" --type py

echo "Checking model deviation calculations in Paddle..."
rg -A 5 "calc_model_devi" --type py

Length of output: 16006

---

Script:

#!/bin/bash
# Check if model deviation is implemented in Paddle backend
echo "Checking model deviation implementation in Paddle backend..."
rg -l "model_devi" deepmd/pd/

# Check DPA2 descriptor implementation in Paddle
echo -e "\nChecking DPA2 descriptor implementation in Paddle..."
rg -l "class.*DPA2" deepmd/pd/model/descriptor/dpa2.py

# Compare DPA2 implementations across backends
echo -e "\nComparing DPA2 implementations across backends..."
for f in deepmd/{pd,pt,jax}/model/descriptor/dpa2.py; do
  if [ -f "$f" ]; then
    echo "=== $f ==="
    head -n 20 "$f"
  fi
done

Length of output: 1855

---

Script:

#!/bin/bash
# Check model deviation implementation in Paddle backend's inference module
echo "Checking for model deviation in Paddle inference..."
fd -e py infer deepmd/pd/ -x cat {}

# Check if there are any model deviation tests for Paddle
echo -e "\nChecking for model deviation tests in Paddle..."
fd -e py "test.*model.*devi" source/tests/pd/ -x cat {}

Length of output: 2272

deepmd/pd/utils/decomp.py (5)

3-7: Consider adding tracking issue for future removal

The comment indicates this file is temporary until Paddle implements these primitives. Consider creating a tracking issue and adding its reference here to ensure this technical debt is properly tracked.

---

53-64: Standardize reduction mode mappings

The reduction modes are mapped inconsistently:

• "sum" -> "add"
• "prod" -> "mul"

Consider standardizing these mappings to match either the input mode names or the underlying Paddle operation names.

---

115-118: Consider caching nonzero indices for repeated operations

The nonzero operation can be expensive for large tensors. If this function is called repeatedly with the same mask, consider caching the computed indices.

+# If this function is called frequently with the same mask, consider adding:
+# @functools.lru_cache(maxsize=128)
+def _compute_mask_coords(mask: paddle.Tensor) -> paddle.Tensor:
+    return paddle.concat(paddle.nonzero(mask, as_tuple=True), axis=1)

---

130-134: Consider using paddle operations in dynamic mode

Instead of using np.prod, consider using paddle's native operations for consistency and potential performance benefits on GPU:

 def numel(x: paddle.Tensor) -> int:
     if paddle.in_dynamic_mode():
-        return np.prod(x.shape)
+        return paddle.prod(paddle.to_tensor(x.shape))
 
     return paddle.numel(x)

---

137-137: Fix typo in comment

"convinience" should be "convenience"

deepmd/pt/model/descriptor/__init__.py (1)

46-58: LGTM! Nice alphabetical ordering.

The __all__ list is now properly alphabetized, which improves readability and maintainability. All exported symbols have corresponding imports.

deepmd/common.py (1)

65-69: LGTM! Consider documenting type hints usage.

The type exports are correctly added to support static type checking. Consider adding a comment or documentation section explaining how to use these type hints in downstream code.

deepmd/pd/train/wrapper.py (4)

30-34: Fix incomplete docstring.

The docstring references a non-existent config parameter and is missing documentation for all actual parameters.

Update the docstring to:

-        """Construct a DeePMD model wrapper.
-
-        Args:
-        - config: The Dict-like configuration with training options.
-        """
+        """Construct a DeePMD model wrapper.
+
+        Args:
+            model: Single model or dictionary of models for multi-task scenarios
+            loss: Optional single loss or dictionary of losses
+            model_params: Optional model parameters dictionary
+            shared_links: Optional parameter sharing configuration
+        """

---

23-29: Add type hints for model_params and shared_links.

Missing type hints could lead to runtime errors if incorrect types are passed.

     def __init__(
         self,
         model: paddle.nn.Layer | dict,
         loss: paddle.nn.Layer | dict = None,
-        model_params=None,
-        shared_links=None,
+        model_params: dict | None = None,
+        shared_links: dict | None = None,
     ) -> None:

---

103-107: Improve error handling for has_spin attribute.

The getattr call could silently ignore missing has_spin implementations.

-        has_spin = getattr(self.model[task_key], "has_spin", False)
+        has_spin = getattr(self.model[task_key], "has_spin", None)
+        if has_spin is None:
+            log.warning(f"Model for task {task_key} does not implement has_spin")
+            has_spin = False
         if callable(has_spin):
             has_spin = has_spin()

---

142-145: Add error handling for extra state.

The method should validate the extra state dictionary structure.

     def set_extra_state(self, extra_state: dict):
+        if not isinstance(extra_state, dict):
+            raise TypeError("extra_state must be a dictionary")
+        if "model_params" not in extra_state or "train_infos" not in extra_state:
+            raise KeyError("extra_state must contain 'model_params' and 'train_infos' keys")
         self.model_params = extra_state["model_params"]
         self.train_infos = extra_state["train_infos"]
         return None

deepmd/pd/model/network/network.py (2)

205-229: Add input validation to forward method.

The forward method lacks input validation for the device parameter, which could lead to runtime errors.

Consider adding validation:

 def forward(self, device: str):
+    if not isinstance(device, (str, paddle.base.libpaddle.Place)):
+        raise TypeError(f"Expected device to be str or Place, got {type(device)}")
     if not self.use_econf_tebd:
         embed = self.embedding_net(
             paddle.eye(self.ntypes, dtype=self.prec).to(device=device)
         )

---

296-321: Enhance error handling in deserialize method.

The deserialize method could benefit from more robust error handling and validation.

Consider adding more comprehensive error handling:

 @classmethod
 def deserialize(cls, data: dict):
+    if not isinstance(data, dict):
+        raise TypeError(f"Expected dict, got {type(data)}")
+    if not data:
+        raise ValueError("Empty data dictionary provided")
+
     data = data.copy()
     check_version_compatibility(data.pop("@version", 1), 2, 1)
     data_cls = data.pop("@class")
     assert data_cls == "TypeEmbedNet", f"Invalid class {data_cls}"
 
+    required_keys = {"ntypes", "neuron", "activation_function"}
+    missing_keys = required_keys - set(data.keys())
+    if missing_keys:
+        raise ValueError(f"Missing required keys: {missing_keys}")
+
     embedding_net = EmbeddingNet.deserialize(data.pop("embedding"))
     # compat with version 1
     if "use_tebd_bias" not in data:
         data["use_tebd_bias"] = True
     type_embedding_net = cls(**data)
     type_embedding_net.embedding_net = embedding_net
     return type_embedding_net

deepmd/pd/model/task/invar_fitting.py (1)

176-176: Fix typo in docstring

There's a typo in the word "alculate" in the docstring.

-    """Based on embedding net output, alculate total energy.
+    """Based on embedding net output, calculate total energy.

deepmd/pd/model/model/__init__.py (2)

75-95: Simplify nested loops and improve error messages.

The nested loops and error handling could be more elegant and informative.

Consider this improvement:

 def _convert_preset_out_bias_to_array(preset_out_bias, type_map):
     if preset_out_bias is not None:
         for kk in preset_out_bias:
             if len(preset_out_bias[kk]) != len(type_map):
                 raise ValueError(
-                    "length of the preset_out_bias should be the same as the type_map"
+                    f"Length mismatch for '{kk}': expected {len(type_map)}, got {len(preset_out_bias[kk])}"
                 )
-            for jj in range(len(preset_out_bias[kk])):
-                if preset_out_bias[kk][jj] is not None:
-                    if isinstance(preset_out_bias[kk][jj], list):
-                        bb = preset_out_bias[kk][jj]
-                    elif _can_be_converted_to_float(preset_out_bias[kk][jj]):
-                        bb = [float(preset_out_bias[kk][jj])]
-                    else:
-                        raise ValueError(
-                            f"unsupported type/value of the {jj}th element of "
-                            f"preset_out_bias['{kk}'] "
-                            f"{type(preset_out_bias[kk][jj])}"
-                        )
-                    preset_out_bias[kk][jj] = np.array(bb)
+            preset_out_bias[kk] = [
+                np.array(v if isinstance(v, list) else [float(v)])
+                if v is not None and (_can_be_converted_to_float(v) or isinstance(v, list))
+                else v
+                for v in preset_out_bias[kk]
+            ]
     return preset_out_bias

---

129-134: Add type hints and documentation.

The function would benefit from type hints and a proper docstring.

Consider this improvement:

-def get_model(model_params):
+def get_model(model_params: dict) -> BaseModel:
+    """Factory method to create model instances.
+    
+    Args:
+        model_params (dict): Model configuration parameters
+        
+    Returns:
+        BaseModel: Created model instance based on the specified type
+        
+    Raises:
+        ValueError: If model_params is invalid
+    """
+    if not isinstance(model_params, dict):
+        raise ValueError("model_params must be a dictionary")
+
     model_type = model_params.get("type", "standard")
     if model_type == "standard":
         return get_standard_model(model_params)
     else:
         return BaseModel.get_class_by_type(model_type).get_model(model_params)

deepmd/pd/model/descriptor/se_atten.py (4)

375-379: Simplify Conditional Assignment Using Ternary Operator

The conditional assignment to sampled can be simplified using a ternary operator for conciseness.

Apply this change:

-            if callable(merged):
-                sampled = merged()
-            else:
-                sampled = merged
+            sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.0)

375-379: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

(SIM108)

---

477-477: Remove Assignment to Unused Variable nall

The variable nall is assigned but never used, which may cause confusion. Consider removing it if it's not needed.

Apply this change:

-            nall = extended_coord.reshape([nb, -1, 3]).shape[1]

🧰 Tools

🪛 Ruff (0.8.0)

477-477: Local variable nall is assigned to but never used

Remove assignment to unused variable nall

(F841)

---

467-470: Clarify Return Types in the Docstring

In the forward method's docstring (lines 467-470), the description of the return values could be more precise. Consider specifying the types and shapes of the returned tensors for better clarity.

Update the docstring as follows:

             Returns
             -------
             result : paddle.Tensor
-                The descriptor. shape: nf x nloc x (ng x axis_neuron)
+                The descriptor tensor with shape: [nf, nloc, ng * axis_neuron]
             g2 : paddle.Tensor
-                The rotationally invariant pair-partical representation.
-                shape: nf x nloc x nnei x ng
+                Rotationally invariant pair-particle representation with shape: [nf, nloc, nnei, ng]
             h2 : paddle.Tensor
-                The rotationally equivariant pair-partical representation.
-                shape: nf x nloc x nnei x 3
+                Rotationally equivariant pair-particle representation with shape: [nf, nloc, nnei, 3]
             gr : paddle.Tensor
-                The rotationally equivariant and permutationally invariant single particle
-                representation. shape: nf x nloc x ng x 3
+                Rotationally equivariant and permutationally invariant single-particle representation with shape: [nf, nloc, ng, 3]
             sw : paddle.Tensor
-                The smooth switch function. shape: nf x nloc x nnei
+                Smooth switching function with shape: [nf, nloc, nnei]

---

587-588: Implement Compression Feature or Remove Unimplemented Code

The method raises a NotImplementedError at lines 587-588, indicating that compression is not implemented yet. If compression is a planned feature, consider implementing it or providing a more descriptive error message. If not, consider removing this code block to avoid runtime errors.

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py (1)

438-447: Consider improving the compression method implementation.

For consistency with set_case_embd, consider:

1. Raising NotImplementedError instead of silently passing
2. Expanding the docstring to explain why compression isn't supported for this model

     def enable_compression(
         self,
         min_nbor_dist: float,
         table_extrapolate: float = 5,
         table_stride_1: float = 0.01,
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Pairtab model does not support compression."""
-        pass
+        """
+        The PairTab model does not support compression as it already uses
+        a pre-computed tabulation for efficiency.
+        """
+        raise NotImplementedError(
+            "Compression not supported for PairTabAtomicModel as it uses pre-computed tabulation!"
+        )

deepmd/pd/model/descriptor/dpa2.py (5)

499-499: Rename unused loop variable ii to _ for clarity

In the loop at line 499, the variable ii is not used within the loop body. Renaming it to _ indicates that the variable is intentionally unused.

Apply this diff to rename the unused variable:

-for ii, descrpt in enumerate(descrpt_list):
+for _, descrpt in enumerate(descrpt_list):
     descrpt.compute_input_stats(merged, path)

🧰 Tools

🪛 Ruff (0.8.0)

499-499: Loop control variable ii not used within loop body

Rename unused ii to _ii

(B007)

---

653-653: Remove unused variable env_mat in repinit deserialization

The variable env_mat is assigned but not used afterward in the deserialization of repinit. Removing this assignment will clean up the code.

Apply this diff to remove the unused assignment:

-    env_mat = repinit_variable.pop("env_mat")

---

668-668: Remove unused variable env_mat in repinit_three_body deserialization

The variable env_mat is assigned but not used afterward in the deserialization of repinit_three_body. Removing this assignment will tidy up the code.

Apply this diff to remove the unused assignment:

-    env_mat = repinit_three_body_variable.pop("env_mat")

---

684-684: Remove unused variable env_mat in repformers deserialization

The variable env_mat is assigned but not used afterward in the deserialization of repformers. Removing this assignment will improve code clarity.

Apply this diff to remove the unused assignment:

-    env_mat = repformers_variable.pop("env_mat")

🧰 Tools

🪛 Ruff (0.8.0)

684-684: Local variable env_mat is assigned to but never used

Remove assignment to unused variable env_mat

(F841)

---

873-897: Implement the enable_compression method or remove the placeholder

The enable_compression method currently raises NotImplementedError. If compression is a planned feature, consider implementing this method. Otherwise, remove the placeholder to avoid confusion and potential runtime errors.

Would you like assistance in implementing the enable_compression method or should we create a GitHub issue to track this task?

deepmd/pd/model/task/fitting.py (4)

419-448: Improve error messages and add performance optimization in _forward_common

The error messages could be more descriptive, and there's an opportunity for performance optimization in tensor operations.

Consider these improvements:

 def _forward_common(self, descriptor: paddle.Tensor, atype: paddle.Tensor, ...):
+    # Validate input shapes early
+    if descriptor.ndim != 3:
+        raise ValueError(
+            f"descriptor must be 3D tensor, got shape {descriptor.shape}"
+        )
+    if atype.ndim != 2:
+        raise ValueError(f"atype must be 2D tensor, got shape {atype.shape}")
+
     # cast the input to internal precsion
-    xx = descriptor.to(self.prec)
-    fparam = fparam.to(self.prec) if fparam is not None else None
-    aparam = aparam.to(self.prec) if aparam is not None else None
+    # Batch convert tensors to improve performance
+    tensors_to_convert = [t for t in [descriptor, fparam, aparam] if t is not None]
+    converted_tensors = [t.to(self.prec) for t in tensors_to_convert]
+    xx, fparam, aparam = (
+        converted_tensors[0],
+        converted_tensors[1] if fparam is not None else None,
+        converted_tensors[2] if aparam is not None else None,
+    )

---

328-331: Optimize dictionary key lookup

The key lookup in the deserialize method can be optimized by removing the .keys() call.

-        for kk in variables.keys():
+        for kk in variables:
             obj[kk] = to_paddle_tensor(variables[kk])

🧰 Tools

🪛 Ruff (0.8.0)

328-328: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

281-320: Document version compatibility in serialize method

The serialize method includes version information but lacks documentation about version compatibility and breaking changes.

Add docstring to explain version compatibility:

 def serialize(self) -> dict:
-    """Serialize the fitting to dict."""
+    """Serialize the fitting to dict.
+    
+    Returns
+    -------
+    dict
+        A dictionary containing the serialized data.
+        Version compatibility:
+        - Version 1: Initial version
+        - Version 2: Added support for mixed types
+        - Version 3: Added support for case embeddings
+    """

---

1-548: Overall well-structured implementation with room for improvements

The code implements a complex fitting network with good architecture and separation of concerns. While the core functionality is solid, consider these general improvements:

1. Add comprehensive docstrings for all public methods
2. Implement logging for better debugging
3. Add unit tests for edge cases
4. Consider adding type hints for internal methods

Would you like assistance in implementing any of these improvements?

🧰 Tools

🪛 Ruff (0.8.0)

132-132: Do not use mutable data structures for argument defaults

Replace with None; initialize within function

(B006)

---

143-143: Do not use mutable data structures for argument defaults

Replace with None; initialize within function

(B006)

---

251-251: Do not use mutable data structures for argument defaults

Replace with None; initialize within function

(B006)

---

328-328: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

deepmd/pd/model/task/ener.py (1)

34-87: Implementation aligns with DPA-2 support objectives

The EnergyFittingNet implementation follows the established patterns of the fitting framework while adding support for case embeddings through dim_case_embd. This aligns well with the PR's objective of implementing DPA-2 support in the Paddle backend.

A few architectural considerations:

1. The class properly extends InvarFitting and maintains the fitting type registration pattern
2. Case embedding support is consistently implemented through constructor parameters
3. Serialization methods maintain proper versioning for future compatibility

🧰 Tools

🪛 Ruff (0.8.0)

40-40: Do not use mutable data structures for argument defaults

Replace with None; initialize within function

(B006)

deepmd/pd/model/descriptor/repformer_layer.py (5)

41-82: Consider adding parameter validation for _scale and _dim

The function handles initialization modes well, but should validate that _scale is non-negative and _dim is positive to prevent potential issues.

def get_residual(
    _dim: int,
    _scale: float,
    _mode: str = "norm",
    trainable: bool = True,
    precision: str = "float64",
    seed: Optional[Union[int, list[int]]] = None,
) -> paddle.Tensor:
+    if _dim <= 0:
+        raise ValueError("_dim must be positive")
+    if _scale < 0:
+        raise ValueError("_scale must be non-negative")
    random_generator = get_generator(seed)

---

85-116: Consider using paddle.gather for better performance

The current implementation uses take_along_axis which is less efficient than gather for this use case.

-    gg1 = paddle.take_along_axis(g1_ext, axis=1, indices=index)
+    gg1 = paddle.gather(g1_ext, index, axis=1)

---

967-967: Remove unused variable assignment

The variable ng2 is assigned but never used in this context.

-    nb, nloc, nnei, _ = g2.shape
-    ng2 = g2.shape[-1]
+    nb, nloc, nnei, _ = g2.shape

🧰 Tools

🪛 Ruff (0.8.0)

967-967: Local variable ng2 is assigned to but never used

Remove assignment to unused variable ng2

(F841)

🪛 GitHub Check: CodeQL

[notice] 967-967: Unused local variable
Variable ng2 is not used.

---

1157-1160: Simplify conditional assignment using ternary operator

The if-else block can be simplified using a ternary operator.

-        if cal_gg1:
-            gg1 = _make_nei_g1(g1_ext, nlist)
-        else:
-            gg1 = None
+        gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None

🧰 Tools

🪛 Ruff (0.8.0)

1157-1160: Use ternary operator gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None instead of if-else-block

Replace if-else-block with gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None

(SIM108)

---

1306-1408: Add validation for required fields in serialization

The serialization method should validate that all required fields are present before creating the dictionary.

def serialize(self) -> dict:
+    required_fields = ['rcut', 'rcut_smth', 'sel', 'ntypes', 'g1_dim', 'g2_dim']
+    for field in required_fields:
+        if not hasattr(self, field):
+            raise ValueError(f"Required field {field} is missing")
    data = {
        "@class": "RepformerLayer",
        "@version": 2,
        # ... rest of the serialization
    }

deepmd/pd/utils/region.py (1)

89-91: Remove or address commented-out code referencing undefined _to_face_distance

The commented-out code references an undefined function _to_face_distance, which may cause confusion or indicate incomplete code.

Consider removing the commented code or implementing _to_face_distance if needed. Apply this diff:

-# b_to_face_distance = paddle.vmap(
-#   _to_face_distance, in_dims=(0), out_dims=(0))

deepmd/pd/utils/dataloader.py (4)

285-290: Simplify nested if statements to improve readability

The nested if statements from lines 285 to 290 can be combined into a single if statement using logical and operators. This will enhance code readability by reducing indentation levels.

Apply this diff to combine the if statements:

-        if self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2)):
-            if time.time() - self.start_time > 5 * 60:
-                if (
-                    self.warning_time is None
-                    or time.time() - self.warning_time > 15 * 60
-                ):
+        if (
+            self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2))
+            and time.time() - self.start_time > 5 * 60
+            and (
+                self.warning_time is None
+                or time.time() - self.warning_time > 15 * 60
+            )
+        ):

🧰 Tools

🪛 Ruff (0.8.0)

285-290: Use a single if statement instead of nested if statements

(SIM102)

---

286-290: Use a single if statement instead of nested if statements

Combine if statements using and

(SIM102)

---

310-310: Iterate directly over the dictionary for efficiency

Instead of using for key in example.keys():, you can iterate directly over the dictionary with for key in example:. This is more Pythonic and slightly more efficient.

Apply this diff to modify the iteration:

-        for key in example.keys():
+        for key in example:

🧰 Tools

🪛 Ruff (0.8.0)

310-310: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

318-319: Add a comment to explain skipping the 'type' key

In the collate_batch function, the code skips processing the 'type' key. Adding a comment to explain the reasoning will improve code clarity and maintainability.

Apply this diff to add a comment:

         elif key == "type":
+            # Skip 'type' key as it does not require collation
             continue

---

98-108: Rename processes variable to threads for clarity

Since multiprocessing.dummy.Pool uses threads rather than processes, renaming the processes variable to threads will prevent confusion and enhance code readability.

Apply this diff to rename the variable:

-        MAX_PROCESSES_NUM = 4
-        processes = min(
+        MAX_THREADS_NUM = 4
+        threads = min(
             os.cpu_count()
             // (
                 dist.get_world_size()
                 if dist.is_available() and dist.is_initialized()
                 else 1
             ),
-            MAX_PROCESSES_NUM,
+            MAX_THREADS_NUM,
         )
-        with Pool(processes) as pool:
+        with Pool(threads) as pool:

deepmd/pt/model/task/fitting.py (1)

215-215: Consider removing the commented-out code for clarity.

The commented line initializing case_embd is not needed and can be removed to keep the code clean.

deepmd/pd/utils/env.py (2)

72-76: Remove unused variables and commented-out code.

The variables inter_nthreads and intra_nthreads are assigned but not used. Additionally, there is commented-out code that appears to be inactive. Removing unused code improves maintainability and readability.

Apply this diff to clean up the code:

 set_default_nthreads()
-inter_nthreads, intra_nthreads = get_default_nthreads()
-# if inter_nthreads > 0:  # the behavior of 0 is not documented
-#     os.environ['OMP_NUM_THREADS'] = str(inter_nthreads)
-# if intra_nthreads > 0:
-#     os.environ['CPU_NUM'] = str(intra_nthreads)

If you intend to use these variables or code in the future, consider adding a comment explaining their purpose or implement them accordingly.

🧰 Tools

🪛 GitHub Check: CodeQL

[notice] 72-72: Unused global variable
The global variable 'inter_nthreads' is not used.

---

[notice] 72-72: Unused global variable
The global variable 'intra_nthreads' is not used.

---

[notice] 75-76: Commented-out code
This comment appears to contain commented-out code.

---

15-16: Remove unused import get_default_nthreads.

After removing the unused variables and code, the function get_default_nthreads is no longer used in this module. Consider removing it from the imports to clean up the code.

Apply this diff to remove the unused import:

 from deepmd.env import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
-    get_default_nthreads,
     set_default_nthreads,
 )

deepmd/pd/model/atomic_model/dp_atomic_model.py (1)

49-54: Consider adding type hints for descriptor and fitting parameters

Adding type hints to the descriptor and fitting parameters in the __init__ method enhances code readability and helps with static type checking. It also makes the expected types clear to other developers.

 def __init__(
-    self,
-    descriptor,
-    fitting,
-    type_map: list[str],
-    **kwargs,
+) -> None:
+    self,
+    descriptor: BaseDescriptor,
+    fitting: BaseFitting,
+    type_map: list[str],
+    **kwargs,
 ) -> None:

deepmd/pd/model/descriptor/env_mat.py (1)

33-34: Improve Variable Naming for Clarity

Variables like t0, t1, and _env_mat_se_a are not descriptive, which can affect code readability and maintainability. Renaming them to more meaningful names would make the code clearer.

Example of improved variable names:

 # Calculations in _make_env_mat
-    t0 = 1 / (length + protection)
-    t1 = diff / (length + protection) ** 2
+    inverse_length = 1 / (length + protection)
+    normalized_diff = diff / (length + protection) ** 2
     # ...
     if radial_only:
-        env_mat = t0 * weight
+        env_mat = inverse_length * weight
     else:
-        env_mat = paddle.concat([t0.astype(t1.dtype), t1], axis=-1) * weight
+        env_mat = paddle.concat([inverse_length.astype(normalized_diff.dtype), normalized_diff], axis=-1) * weight
     # ...

# Normalization in prod_env_mat
-    env_mat_se_a = (_env_mat_se_a - t_avg) / t_std
+    normalized_env_mat = (env_mat - t_avg) / t_std

Also applies to: 81-81

deepmd/pd/entrypoints/main.py (2)

116-118: Remove unused variable f

The variable f is assigned but never used. Since you're creating the file without needing to interact with it, you can omit assigning it to f.

Apply this diff to remove the unused variable:

            if stat_file_path_single.endswith((".h5", ".hdf5")):
-               with h5py.File(stat_file_path_single, "w") as f:
+               with h5py.File(stat_file_path_single, "w"):
                    pass

🧰 Tools

🪛 Ruff (0.8.0)

118-118: Local variable f is assigned to but never used

Remove assignment to unused variable f

(F841)

---

500-503: Simplify conditional assignment of FLAGS

You can use a ternary operator to simplify the assignment of FLAGS.

Apply this diff to simplify the code:

-        if not isinstance(args, argparse.Namespace):
-            FLAGS = parse_args(args=args)
-        else:
-            FLAGS = args
+        FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args

🧰 Tools

🪛 Ruff (0.8.0)

500-503: Use ternary operator FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args instead of if-else-block

Replace if-else-block with FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args

(SIM108)

deepmd/pd/model/descriptor/se_a.py (3)

47-52: Simplify import handling using contextlib.suppress

Replace the try-except-pass block with contextlib.suppress for cleaner code.

Apply this diff:

+    from contextlib import suppress
     # ...
-    try:
-        from typing import (
-            Final,
-        )
-    except ImportError:
-        pass
+    with suppress(ImportError):
+        from typing import Final

🧰 Tools

🪛 Ruff (0.8.0)

47-52: Use contextlib.suppress(ImportError) instead of try-except-pass

Replace with contextlib.suppress(ImportError)

(SIM105)

🪛 GitHub Check: CodeQL

[notice] 50-50: Empty except
'except' clause does nothing but pass and there is no explanatory comment.

---

361-361: Remove unused variable env_mat

The variable env_mat is assigned but never used. You can remove this assignment to clean up the code.

Apply this diff:

             embeddings = data.pop("embeddings")
-            env_mat = data.pop("env_mat")
             obj = cls(**data)

🧰 Tools

🪛 Ruff (0.8.0)

361-361: Local variable env_mat is assigned to but never used

Remove assignment to unused variable env_mat

(F841)

---

645-650: Remove redundant assignment to ti

The variable ti is assigned a value but then immediately reassigned before being used. Remove the initial assignment to ti to simplify the code.

Apply this diff:

                if self.type_one_side:
                    ii = embedding_idx
-                   ti = -1
                    # paddle.jit is not happy with slice(None)
                    # ti_mask = paddle.ones(nfnl, dtype=paddle.bool, device=dmatrix.place)
                    # applying a mask seems to cause performance degradation
                    ti_mask = None
                else:
                    # ti: center atom type, ii: neighbor type...
                    ii = embedding_idx // self.ntypes
                    ti = embedding_idx % self.ntypes

deepmd/pd/train/training.py (4)

290-290: Simplify dictionary get calls by omitting redundant None

When using dict.get(key, default), you can omit the default if it's None, as it's the default value.

Apply this diff:

290,392
-                    if config.get("learning_rate_dict", None) is not None:
+                    if config.get("learning_rate_dict") is not None:

Also applies to: 392-392

🧰 Tools

🪛 Ruff (0.8.0)

290-290: Use config.get("learning_rate_dict") instead of config.get("learning_rate_dict", None)

Replace config.get("learning_rate_dict", None) with config.get("learning_rate_dict")

(SIM910)

---

493-493: Simplify key existence check in dictionary

You can simplify the key existence check by removing .keys().

Apply this diff:

-                        if key in _random_state_dict.keys():
+                        if key in _random_state_dict:

🧰 Tools

🪛 Ruff (0.8.0)

493-493: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

863-867: Remove unused variable module

The variable module is assigned but never used. You can remove this assignment to clean up the code.

Apply this diff:

-                    module = (
-                        self.wrapper.module
-                        if dist.is_available() and dist.is_initialized()
-                        else self.wrapper
-                    )
                    self.save_model(self.latest_model, lr=cur_lr, step=_step_id)

🧰 Tools

🪛 Ruff (0.8.0)

863-863: Local variable module is assigned to but never used

Remove assignment to unused variable module

(F841)

---

1048-1048: Simplify iteration over dictionary

When iterating over a dictionary, you can iterate directly over it without calling .keys().

Apply this diff:

-            for key in batch_data.keys():
+            for key in batch_data:

🧰 Tools

🪛 Ruff (0.8.0)

1048-1048: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

deepmd/pd/utils/utils.py (3)

35-57: Consider using an enum for activation function types and optimize string comparisons.

The activation function implementation could be improved in terms of maintainability and performance.

Consider these improvements:

+from enum import Enum, auto
+
+class ActivationType(Enum):
+    RELU = auto()
+    GELU = auto()
+    TANH = auto()
+    RELU6 = auto()
+    SOFTPLUS = auto()
+    SIGMOID = auto()
+    LINEAR = auto()
+
 class ActivationFn(paddle.nn.Layer):
     def __init__(self, activation: str | None):
         super().__init__()
-        self.activation: str = activation if activation is not None else "linear"
+        self.activation: str = (activation or "linear").lower()
 
     def forward(self, x: paddle.Tensor) -> paddle.Tensor:
         """Returns the tensor after applying activation function corresponding to `activation`."""
-        if self.activation.lower() == "relu":
+        if self.activation == "relu":
             return F.relu(x)
-        elif self.activation.lower() == "gelu" or self.activation.lower() == "gelu_tf":
+        elif self.activation in ("gelu", "gelu_tf"):
             return F.gelu(x, approximate=True)
-        elif self.activation.lower() == "tanh":
+        elif self.activation == "tanh":
             return paddle.tanh(x)
-        elif self.activation.lower() == "relu6":
+        elif self.activation == "relu6":
             return F.relu6(x)
-        elif self.activation.lower() == "softplus":
+        elif self.activation == "softplus":
             return F.softplus(x)
-        elif self.activation.lower() == "sigmoid":
+        elif self.activation == "sigmoid":
             return F.sigmoid(x)
-        elif self.activation.lower() == "linear" or self.activation.lower() == "none":
+        elif self.activation in ("linear", "none"):
             return x
         else:
-            raise RuntimeError(f"activation function {self.activation} not supported")
+            supported = ["relu", "gelu", "tanh", "relu6", "softplus", "sigmoid", "linear", "none"]
+            raise RuntimeError(f"Activation function '{self.activation}' not supported. Supported functions: {supported}")

🧰 Tools

🪛 GitHub Check: CodeQL

[notice] 55-55: Statement has no effect
This statement has no effect.

---

116-127: Simplify device management logic and add device validation.

The device management function could be more concise and robust.

Consider this improvement:

 def dict_to_device(sample_dict):
+    valid_devices = {"cpu", "gpu", "xpu"}  # Add other valid devices as needed
+    if not any(DEVICE.startswith(dev) for dev in valid_devices):
+        raise ValueError(f"Invalid device: {DEVICE}")
+
     for key in sample_dict:
-        if isinstance(sample_dict[key], list):
-            sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]]
-        if isinstance(sample_dict[key], np.float32):
-            sample_dict[key] = (
-                paddle.ones(1, dtype=paddle.float32).to(device=DEVICE)
-                * sample_dict[key]
-            )
-        else:
-            if sample_dict[key] is not None:
-                sample_dict[key] = sample_dict[key].to(DEVICE)
+        value = sample_dict[key]
+        if value is None:
+            continue
+        if isinstance(value, list):
+            sample_dict[key] = [item.to(DEVICE) for item in value]
+        elif isinstance(value, np.float32):
+            sample_dict[key] = paddle.ones(1, dtype=paddle.float32).to(device=DEVICE) * value
+        else:
+            sample_dict[key] = value.to(DEVICE)

---

188-198: Enhance error handling in profiling context manager.

The profiling context manager could benefit from better error handling and logging.

Consider this improvement:

+import logging
+
+logger = logging.getLogger(__name__)
+
 @contextmanager
 def nvprof_context(enable_profiler: bool, name: str):
     if enable_profiler:
-        core.nvprof_nvtx_push(name)
+        try:
+            core.nvprof_nvtx_push(name)
+            logger.debug(f"Started profiling section: {name}")
+        except Exception as e:
+            logger.error(f"Failed to start profiling section {name}: {str(e)}")
+            raise
 
     try:
         yield
 
     finally:
         if enable_profiler:
-            core.nvprof_nvtx_pop()
+            try:
+                core.nvprof_nvtx_pop()
+                logger.debug(f"Ended profiling section: {name}")
+            except Exception as e:
+                logger.error(f"Failed to end profiling section {name}: {str(e)}")

deepmd/entrypoints/test.py (4)

707-709: Consider using pathlib's stem and suffix methods for file path manipulation.

The current implementation directly concatenates strings for file suffixes. While this works, using pathlib's built-in methods would be more idiomatic and safer.

-                detail_path.with_suffix(f".dos.out.{ii}"),
+                detail_path.with_stem(f"{detail_path.stem}.dos.{ii}").with_suffix(".out"),

---

721-723: Maintain consistency with the previous file naming pattern.

The atomic DOS output follows a similar pattern to the regular DOS output. For consistency, consider applying the same refactoring suggestion.

-                    detail_path.with_suffix(f".ados.out.{ii}"),
+                    detail_path.with_stem(f"{detail_path.stem}.ados.{ii}").with_suffix(".out"),

---

863-865: Maintain consistent file naming pattern across test functions.

The property output follows the same pattern as DOS output. For consistency, consider applying the same refactoring.

-                detail_path.with_suffix(f".property.out.{ii}"),
+                detail_path.with_stem(f"{detail_path.stem}.property.{ii}").with_suffix(".out"),

---

877-879: Ensure consistent file naming pattern for atomic properties.

The atomic property output should follow the same pattern as other atomic outputs.

-                    detail_path.with_suffix(f".aproperty.out.{ii}"),
+                    detail_path.with_stem(f"{detail_path.stem}.aproperty.{ii}").with_suffix(".out"),

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Description: Find where jax_export is being imported/used in the codebase
# This helps understand the impact of exposing this functionality

# Search for imports of jax_export through this module
rg -l "from deepmd.jax.env import.*jax_export"

# Search for direct usage of jax_export
rg "jax_export\." 

#!/bin/bash
# Description: Find all usages of calculate_gain with leaky_relu to ensure compatibility
# with the new type checking logic

# Search for direct calls to calculate_gain with leaky_relu
rg -l "calculate_gain.*leaky_relu" 

# Search for kaiming initialization calls that might use leaky_relu (default nonlinearity)
rg -l "kaiming_(uniform|normal)_"

#!/bin/bash
# Let's examine the actual usage of calculate_gain and kaiming initialization
# to ensure the type safety changes don't break existing code

# Check calculate_gain usage in mlp.py files
echo "=== MLP Usage in PyTorch implementation ==="
rg "calculate_gain|kaiming_" deepmd/pt/model/network/mlp.py -B 2 -A 2

echo -e "\n=== Network Usage in PyTorch implementation ==="
rg "calculate_gain|kaiming_" deepmd/pt/model/network/network.py -B 2 -A 2

echo -e "\n=== MLP Usage in PaddlePaddle implementation ==="
rg "calculate_gain|kaiming_" deepmd/pd/model/network/mlp.py -B 2 -A 2

    "BaseAtomicModel",
    "DPAtomicModel",
    "DPDOSAtomicModel",
    "DPDipoleAtomicModel",
    "DPEnergyAtomicModel",
    "DPPolarAtomicModel",
    "DPPropertyAtomicModel",
    "DPZBLLinearEnergyAtomicModel",
    "LinearEnergyAtomicModel",
    "PairTabAtomicModel",

    "DOSModel",
    "DipoleModel",
    "EnerModel",
    "GlobalPolarModel",
    "LinearEnergyModel",
    "PairTabModel",
    "PolarModel",

#!/bin/bash
# Description: Check for similar ordering patterns in other __init__.py files
# Look for other __init__.py files with __all__ declarations to verify consistency

rg -U "__all__\s*=\s*\[[\s\S]*?\]" --type py --glob "*__init__.py"

#!/bin/bash
# Description: Check documentation and usage of DPModelCommon

# Check if class has proper docstrings
echo "Checking DPModelCommon documentation..."
ast-grep --pattern 'class DPModelCommon {
  $$$
}'

# Check for existing usage patterns
echo "Checking current usage patterns..."
rg -l "DPModelCommon" --type python

# Check if this might break existing imports
echo "Checking potential import conflicts..."
rg -l "from deepmd.dpmodel.model import \*" --type python

#!/bin/bash
# Let's try a different approach without file type filters

# Check class definition and documentation
echo "Checking DPModelCommon class definition..."
rg -A 10 "class DPModelCommon"

# Check current usage patterns
echo "Checking current usage patterns..."
rg "DPModelCommon"

# Check potential import conflicts
echo "Checking import patterns..."
rg "from deepmd.dpmodel.model import"

#!/bin/bash
# Description: Check for version-related documentation and consistency

# Check for version numbers in documentation
echo "Checking for version documentation..."
rg -i "version.*3" docs/ || echo "No version 3 documentation found"

# Check for other version checks to ensure consistency
echo "Checking for other version checks..."
rg "check_version_compatibility.*2|check_version_compatibility.*3" || echo "No other version checks found"

#!/bin/bash
# Search for usage of the newly exported variables in the Paddle backend
rg "GLOBAL_ENER_FLOAT_PRECISION|GLOBAL_NP_FLOAT_PRECISION|global_float_prec" --type python

#!/bin/bash
# Search for usage of the newly exported variables without file type restriction
rg "GLOBAL_ENER_FLOAT_PRECISION|GLOBAL_NP_FLOAT_PRECISION|SHARED_LIB_MODULE|global_float_prec"

#!/bin/bash
# Search for set_case_embd implementation in atomic models
ast-grep --pattern 'def set_case_embd(self, case_idx: int):'

#!/bin/bash
# Let's try a broader search for set_case_embd method using ripgrep
rg "set_case_embd" -A 3

#!/bin/bash
# Description: Verify consistency between imports and exports in __init__.py

# Extract imported names and __all__ list contents
echo "Checking deepmd/dpmodel/fitting/__init__.py for import/export consistency..."

# Get all imported names (excluding 'from' and 'import' keywords)
imported=$(awk -F'[{,}]' '/^from.*import.*/ {
    # Extract everything between "import (" and ")"
    match($0, /import \((.*)\)/, arr)
    if (arr[1] != "") {
        # Clean up the captured group
        gsub(/[ \n]/, "", arr[1])
        print arr[1]
    }
}' deepmd/dpmodel/fitting/__init__.py)

# Get all exported names from __all__
exported=$(awk '/^__all__/,/]/ {
    # Remove __all__ = [ and ] and clean up
    gsub(/^__all__ = \[|]|"| /, "")
    if (NF > 0) print
}' deepmd/dpmodel/fitting/__init__.py)

# Compare the lists
echo "Imported names:"
echo "$imported"
echo "Exported names:"
echo "$exported"

# Check if any names are in one list but not the other
echo "Checking for mismatches..."

#!/bin/bash
# Description: Verify consistency between imports and exports using grep

echo "Imported names:"
grep -o "from .* import.*(" deepmd/dpmodel/fitting/__init__.py | while read -r line; do
    # Get the next line after each import statement that contains a closing parenthesis
    line_num=$(grep -n "$line" deepmd/dpmodel/fitting/__init__.py | cut -d: -f1)
    sed -n "${line_num},/)/p" deepmd/dpmodel/fitting/__init__.py
done

echo -e "\nExported names:"
grep -A 10 "^__all__" deepmd/dpmodel/fitting/__init__.py

#!/bin/bash
# Description: Check if GeneralFitting class properly handles dim_case_embd
# Look for the parameter in the parent class implementation
rg -A 10 "class GeneralFitting" deepmd/dpmodel/fitting/

#!/bin/bash
# Check the constructor and parameter handling in GeneralFitting
rg -B 5 -A 20 "__init__" deepmd/dpmodel/fitting/general_fitting.py

#!/bin/bash
# Search for usage of dim_case_embd in GeneralFitting implementation
rg "dim_case_embd" deepmd/dpmodel/fitting/general_fitting.py

#!/bin/bash
# Description: Verify the error message formatting by checking for any string formatting issues
# or unclosed f-string brackets in the codebase

# Check for potential f-string formatting issues
rg -U "f['\"].*?['\"]" deepmd/pt/cxx_op.py | while read -r line; do
    if echo "$line" | grep -q "{.*{\\|}.*}"; then
        echo "Potential nested brackets in f-string: $line"
    fi
    if echo "$line" | grep -q "{[^}]*$"; then
        echo "Potential unclosed bracket in f-string: $line"
    fi
done

# Check for consistency in error message formatting
ast-grep --pattern 'raise RuntimeError(f"$_")' deepmd/pt/cxx_op.py

#!/bin/bash
# Let's check the complete context of the error message to ensure proper formatting
rg -U "raise RuntimeError.*?\)" deepmd/pt/cxx_op.py -A 5

#!/bin/bash
# Let's get more context around the error message construction and usage
rg -U "error_message.*?=.*?\n.*?raise RuntimeError" deepmd/pt/cxx_op.py -B 10

#!/bin/bash
# Let's examine the entire file content to understand the error message context
cat deepmd/pt/cxx_op.py

#!/bin/bash
# Description: Check for corresponding changes in Paddle backend
# Look for similar files and DPA-2 related changes in the Paddle backend

# Check for corresponding model initialization files in Paddle backend
fd -t f "__init__.py" "deepmd/pd"

# Look for DPA-2 related files or changes
rg -l "dpa2|DPA-2|DPA2" "deepmd/pd"

#!/bin/bash
# Let's examine the corresponding model/__init__.py in Paddle backend and DPA2 changes

# Check the corresponding model/__init__.py in Paddle backend
cat deepmd/pd/model/model/__init__.py

# Look at DPA2 descriptor implementation
cat deepmd/pd/model/descriptor/dpa2.py

# Check descriptor/__init__.py for DPA2 registration
cat deepmd/pd/model/descriptor/__init__.py

#!/bin/bash
# Search for version migration logic or documentation
rg -A 5 "check_version_compatibility.*2.*3" --type py

#!/bin/bash
# Let's first check for any version compatibility checks in the codebase
rg -A 5 "check_version_compatibility" --type py

# Also search for version-related constants or variables
rg "VERSION|version" --type py

#!/bin/bash
# Let's check the version.py file to understand the version compatibility logic
cat deepmd/utils/version.py

# Also check for any version migration documentation
rg -l "version.*migration|migration.*version" --type md

#!/bin/bash
# Search for tensor shape assumptions in the codebase
rg -A 2 "tensor\.shape" --type python

# Search for Paddle tensor creation patterns
ast-grep --pattern 'paddle.empty($$$)'

#!/bin/bash
# Search for tensor shape handling in initialization functions
rg -A 5 "\.shape" deepmd/pd/model/network/init.py

# Search for tensor operations to understand layout conventions
rg -A 3 "tensor\." deepmd/pd/model/network/init.py

# Look for any framework-specific imports or configurations
rg -l "import paddle" deepmd/pd/model/network/

#!/bin/bash
# Look for tensor layout documentation or usage in tests
rg -B 2 -A 2 "reverse" deepmd/pd/model/network/init.py

# Check if there are any test files that verify initialization
fd -e py test_init deepmd/pd/model/network/

#!/bin/bash
# Description: Search for usage of sharing levels 2 and 3 in the codebase

# Search for direct usage of share_params with levels 2 or 3
rg -A 5 'share_params.*(?:2|3)' 

# Search for configuration files that might specify these sharing levels
rg -A 5 'shared_level.*(?:2|3)'

#!/bin/bash
# Let's try a broader search to find any references to shared_level or parameter sharing

# Search for shared_level in Python files
rg "shared_level" --type py -B 2 -A 2

# Search for share_params method calls or definitions
rg "share_params" --type py -B 2 -A 2

# Search for any test files that might be testing parameter sharing
fd "test.*\.py" --exec rg -l "share"

#!/bin/bash
# Let's check the test files for any usage of sharing levels 2 and 3
rg -A 5 "shared_level.*[23]" ./source/tests/pd/test_multitask.py ./source/tests/pt/test_multitask.py

# Also check configuration examples for any usage of these levels
rg -A 5 "shared_level.*[23]" ./source/tests/common/test_examples.py

#!/bin/bash
# Description: Check the DPA-2 implementation across the codebase

# Check for DPA-2 related files and their contents
echo "Searching for DPA-2 related files..."
fd "dpa2" --type f

# Check for DPA-2 class implementation details
echo -e "\nChecking DPA-2 class implementation..."
ast-grep --pattern 'class DescrptDPA2' -A 10

# Look for DPA-2 tests
echo -e "\nSearching for DPA-2 tests..."
rg "test.*dpa2" --type py

#!/bin/bash
# Check the implementation of DescrptDPA2 in different frameworks
echo "JAX Implementation:"
cat deepmd/jax/descriptor/dpa2.py

echo -e "\nPyTorch Implementation:"
cat deepmd/pt/model/descriptor/dpa2.py

echo -e "\nPaddlePaddle Implementation:"
cat deepmd/pd/model/descriptor/dpa2.py

echo -e "\nTest Implementation Details:"
cat source/tests/common/dpmodel/test_descriptor_dpa2.py
cat source/tests/consistent/descriptor/test_dpa2.py

#!/bin/bash
# Check test files for DPA2
echo "Common Tests:"
cat source/tests/common/dpmodel/test_descriptor_dpa2.py

echo -e "\nConsistency Tests:"
cat source/tests/consistent/descriptor/test_dpa2.py

echo -e "\nPyTorch Tests:"
cat source/tests/pt/model/test_descriptor_dpa2.py
cat source/tests/pt/model/test_compressed_descriptor_dpa2.py

echo -e "\nPaddlePaddle Tests:"
cat source/tests/pd/model/test_descriptor_dpa2.py

#!/bin/bash
# Search for usage of these operations in DPA-2 related code
echo "Searching for usage of decomposed operations..."
rg -l "masked_add_|scatter_reduce" --type py

# Check if there are any similar operations in torch backend for comparison
echo "Checking torch backend implementation..."
rg -l "masked_add|scatter_reduce" --type py | grep "torch"

#!/bin/bash
# Let's examine the actual usage patterns in the identified files
echo "Checking usage in transform_output.py files..."
rg -B2 -A2 "masked_add_|scatter_reduce" deepmd/pd/model/model/transform_output.py deepmd/pt/model/model/transform_output.py

echo "\nChecking test implementations..."
rg -B2 -A2 "masked_add_|scatter_reduce" source/tests/pd/test_decomp.py source/tests/pd/model/test_forward_lower.py

# Check if these operations are used consistently across backends
echo "\nComparing implementations between paddle and torch..."
ast-grep --pattern 'def masked_add_($$$)'
ast-grep --pattern 'def scatter_reduce($$$)'

#!/bin/bash
# Description: Check if all potential sub-model types support set_case_embd
# Look for class definitions that could be used as sub-models

ast-grep --pattern 'class $name($_) {
  $$$
  set_case_embd($$$) {
    $$$
  }
  $$$
}'

#!/bin/bash
# Let's try a different approach to find all classes that might be used as sub-models
# and check their method definitions

# First, let's find where these models are instantiated or imported
rg -A 5 "self\.models\s*=" 

# Also search for set_case_embd method implementations
rg "def set_case_embd"

-    def set_case_embd(self, case_idx: int):
+    def set_case_embd(self, case_idx: int) -> None:
         """
         Set the case embedding of this atomic model by the given case_idx,
         typically concatenated with the output of the descriptor and fed into the fitting net.
         """
-        # only set case_idx for dpmodel
-        self.models[0].set_case_embd(case_idx)
+        if not self.models:
+            raise ValueError("No models available")
+        dp_model = self.models[0]
+        if not isinstance(dp_model, DPAtomicModel):
+            raise ValueError("First model must be a DPAtomicModel")
+        dp_model.set_case_embd(case_idx)

#!/bin/bash
# Description: Check for any potential DPA-2 related classes/functions that might need to be exposed

# Search for DPA-2 related classes/functions in the codebase
echo "Searching for DPA-2 related definitions..."
ast-grep --pattern 'class $name {
  $$$
}' | rg -A 5 'DPA2|dpa2'

# Search for DPA-2 related imports in other files
echo "Checking usage of DPA-2 components..."
rg -l "from deepmd.pt.model.descriptor import" | xargs rg "DescrptDPA2"

#!/bin/bash
# Check for DPA-2 related parameter classes and their usage
echo "Searching for DPA2 parameter classes..."
ast-grep --pattern 'class $name {
  $$$
}' | rg -B 2 -A 5 'DescriptorParamDPA2'

# Check the actual implementation file for DPA2
echo "Checking DPA2 implementation imports..."
fd "descriptor_dpa2.py" --exec cat {}

# Check if there are any utility functions related to DPA2
echo "Checking for DPA2 related utility functions..."
rg -l "def.*dpa2" | xargs cat

#!/bin/bash
# Search for DPA-2 descriptor implementation
echo "Searching for DPA-2 descriptor implementation..."
fd "dpa2.py" --exec cat {}

# Search for DPA-2 related imports and usage
echo "Checking DPA-2 imports and usage..."
rg "class DescrptDPA2"

#!/bin/bash
# Search for Paddle-specific usage of these types
rg --type python "import paddle|from paddle" -B 2 -A 2

# Search for DPA-2 related code that might use these types
rg --type python "dpa2|DPA2|DPA-2" -B 2 -A 2

#!/bin/bash
# Search for Paddle-specific usage without file type restriction
rg "import paddle|from paddle" -B 2 -A 2

# Search for activation functions usage in paddle-related files
rg -g "*.py" "(relu|relu6|softplus|sigmoid|tanh|gelu|gelu_tf|none|linear)" "$(fd -g "*paddle*.py")"

# Search for DPA-2 related code
rg "dpa2|DPA2|DPA-2" -B 2 -A 2

# Look for paddle backend implementation files
fd -g "*paddle*.py"

#!/bin/bash
# Search for activation function usage in Paddle backend
rg -g "*.py" "(relu|relu6|softplus|sigmoid|tanh|gelu|gelu_tf|none|linear)" "deepmd/pd/"

# Search for activation function implementations in Paddle backend
ast-grep --pattern $'def $_(self, $_, $$$) {
  $$$
  paddle.nn.functional.$activation($$$)
  $$$
}' deepmd/pd/

#!/bin/bash
# Search for activation function usage in the codebase
rg --type python "activation.*=|activation.*:" -B 2 -A 2

#!/bin/bash
# Search for activation function usage in the codebase without file type restriction
rg "activation.*=|activation.*:" -B 2 -A 2

#!/bin/bash
# Search for imports or usage of activation functions
rg -i "activation|relu|softplus|sigmoid|tanh|gelu" -B 2 -A 2

#!/bin/bash
# Find Python files that might contain activation-related code
fd -e py -x cat {} | rg -i "activation|relu|softplus|sigmoid|tanh|gelu" -B 2 -A 2

#!/bin/bash
# Use ast-grep to find function definitions or class methods related to activation
ast-grep --pattern 'def $_($$$) {
  $$$
  activation
  $$$
}'

#!/bin/bash
# Check if parameter sharing is implemented in other files or branches
rg -g '!wrapper.py' "share_params.*multitask" .

#!/bin/bash
# Search for any implementation or references to share_params
rg "share_params" .

# Also check for multitask related code
rg "multitask" .

# Look for any parameter sharing patterns in Python files
ast-grep --pattern 'def share_params($$$)'

#!/bin/bash
# Search for version compatibility checks
rg "check_version_compatibility" --type python -B 2 -A 2

#!/bin/bash
# Search for version compatibility checks without file type restriction
rg "check_version_compatibility" -B 2 -A 2

#!/bin/bash
# Find the definition of check_version_compatibility
ast-grep --pattern 'def check_version_compatibility($$$)'

#!/bin/bash
# Find all Python files that might contain version-related constants
rg "VERSION|version" --type-add 'py:*.py' --type py

#!/bin/bash
# Find the definition of check_version_compatibility in version.py
cat deepmd/utils/version.py

#!/bin/bash
# Find any documentation or comments about version compatibility in the codebase
rg -B 2 -A 2 "version.*compatibility|compatibility.*version" --type-add 'py:*.py' --type py

#!/bin/bash
# Description: Check for DPA-2 and Paddle backend integration points

# Look for DPA-2 related configurations and usage
echo "Checking for DPA-2 configurations..."
rg -l "dpa2|DPA-2|DPA2" --type py

# Check for Paddle backend integration
echo "Checking for Paddle backend integration..."
rg -l "paddle|Paddle" --type py

# Look for related model implementations
echo "Checking for related model implementations..."
fd -e py -x rg -l "class.*Model.*paddle" {}

#!/bin/bash
# Let's verify the model integration more specifically

# Check model implementation in pd backend
echo "=== Checking Paddle model implementation ==="
rg "class.*Model" deepmd/pd/model/model/model.py -A 10

# Check DPA2 descriptor implementation
echo -e "\n=== Checking DPA2 descriptor implementation ==="
rg "class" deepmd/pd/model/descriptor/dpa2.py -A 5

# Check model factory and integration points
echo -e "\n=== Checking model factory implementation ==="
rg "def get_model|def get_standard_model" deepmd/pd/model/model/__init__.py -A 10

# Check for any potential integration issues
echo -e "\n=== Checking for potential integration issues ==="
rg "RuntimeError|ValueError" deepmd/pd/model/model/__init__.py -A 2

#!/bin/bash
# Description: Check for version compatibility documentation

# Check README and docs for version compatibility info
rg -i "version.*compatibility|migration" README.md docs/

# Check for other version checks in the codebase
rg "check_version_compatibility.*3" --type py

#!/bin/bash
# Search for version-related changes in git history
git log -n 10 --grep="version" --pretty=format:"%h %s"

# Look for version-related constants or utilities
rg "VERSION|version" --type py -B 2 -A 2

# Check for any test files related to version compatibility
fd "test.*version.*\.py$"

#!/bin/bash
# Check for version compatibility check implementation
rg "check_version_compatibility" --type py -B 2 -A 2

# Look for version-related changes in the model
rg "MODEL_VERSION" --type py -B 2 -A 2

#!/bin/bash
# Check the version compatibility implementation
rg -A 10 "def check_version_compatibility" deepmd/utils/version.py

# Look for version-related changes in git history
git log -n 5 --oneline deepmd/pd/model/task/ener.py

#!/bin/bash
# Get the complete version compatibility implementation
cat deepmd/utils/version.py

# Check for any version migration utilities
rg "convert.*version" --type py

#!/bin/bash
# Description: Verify that DEVICE is correctly assigned based on environment variables and CUDA availability.

python -c "
import os
import paddle
LOCAL_RANK = paddle.distributed.get_rank()
if os.environ.get('DEVICE') == 'cpu' or paddle.device.cuda.device_count() <= 0:
    DEVICE = 'cpu'
else:
    DEVICE = f'gpu:{LOCAL_RANK}'
print(f'DEVICE is set to: {DEVICE}')
"

#!/bin/bash
# Let's analyze the codebase to understand the device selection logic and its usage

# Check how LOCAL_RANK is defined and used
rg "LOCAL_RANK" -A 3 -B 3

# Check for any device-related configurations or settings
rg "DEVICE\s*=" -A 3 -B 3

# Look for paddle device-related code
rg "paddle.device" -A 3 -B 3

# Search for distributed training related code
rg "paddle.distributed" -A 3 -B 3

#!/bin/bash
# Description: Check file naming patterns in test functions
# Expected: Find all file output operations to verify naming consistency

# Search for file suffix patterns in test functions
rg -A 2 'with_suffix\(' deepmd/entrypoints/test.py

# Search for other file output operations that might use different patterns
rg -A 2 '\.out\.' deepmd/entrypoints/test.py